#!/bin/env python3
import ithildin as ith
import numpy as np
import matplotlib.pyplot as plt
import sys
from typing import List
from scipy.interpolate import RectBivariateSpline
from pydiffmap.diffusion_map import DiffusionMap
from pydiffmap import visualization as diff_visualization
import matplotlib

# Courtemanche1998 model, 2 ruimtedimensies, 21 variabelen

data = ith.SimData.from_stem("myokit12/myokit_2")

# Vermijd randeffecten en initialiatie
# TODO: begintijd zou in principe in de log moeten staan
def snij_randen_weg(variabele,begintijd, eindtijd):
    def begin(index):
        if index == 0:
            return begintijd
        return variabele.shape[index]//8
    def eind(index):
        if index == 0:
            return eindtijd
        if variabele.shape[index] < 20:
            return variabele.shape[index]
        return (7 * variabele.shape[index])//8
    return variabele[begin(0):eind(0),begin(1):eind(1),begin(2):eind(2),begin(3):eind(3)]

def snij_randen_weg2(variables,begintijd,eindtijd):
    newdata = dict()
    for key in variables.keys():
        newdata[key] = snij_randen_weg(variables[key], begintijd,eindtijd)
    return newdata

def onderzoek(vars, vt0, t1, aantal_punten, eps, var1, var2, var3, var4,k=64):
    vars = snij_randen_weg2(vars, t0, t1)
    # We bekijken de faseruimte, tijds- en ruimtecoordinaten zijn daarin niet van belang.
    for key in vars.keys():
        vars[key] = vars[key].ravel()
    # Verlaag aantal punten om geheugengebruik en tijdsduur te beperken.
    # replace=True strikt genomen incorrect maar verlaagt geheugengebruik
    # en wegens de hoge hoeveelheid punten weinig belangrijk
    keuzes = np.random.choice(vars['u'].shape[0], aantal_punten,replace=True)
    for key in vars.keys():
        vars[key] = vars[key][keuzes]
    # Nu de randen verwijderd zijn en we punten in een faseruimte hebben, kunnens we proberen
    # diffusion maps te gebruiken.
    dmap = DiffusionMap.from_sklearn(epsilon=eps,n_evecs=4,k=k)
    dmap.fit(np.column_stack((vars[var1], vars[var2], vars[var3], vars[var4])))
    scatter_kwargs = {'c' : dmap.dmap[:,3] }
    diff_visualization.embedding_plot(dmap,dim=3, scatter_kwargs=scatter_kwargs)
    return dmap

# Different variables are in different units, so normalise them
vars = dict(data.vars)
assert(np.std(data.vars['K']) == 0) # it is constant, so don't plot it (the contrentation of K+ ions remains constant?)
del vars['K']
for key in vars.keys(): # K is constant
    vars[key] = vars[key] - np.mean(vars[key])
    vars[key] = vars[key] / np.std(vars[key])

t0 = 1
t1 = 150
aantal_punten = 5000
matplotlib.use("TkAgg")
# Choose random combinations of vour variables.
keys = list(vars.keys())
keys.sort() # determinism

np.random.seed(1234) # determinism
for i in range(0,500):
    which_keys = np.random.choice(keys, 4, replace=False)
    print(which_keys)
    onderzoek(vars, t0, t1, aantal_punten, 20, which_keys[0], which_keys[1], which_keys[2], which_keys[3],k=2048)

# 3Dig?
# ['gateh' 'gatexs' 'gated' 'gateoa'] eps=20/k=2048
# ['gateu' 'gatev' 'gatef' 'gatefCa']
# ['Na' 'gateh' 'gatexr' 'gatef']
# ['gatew' 'gateh' 'gatexr' 'gateua']
# ['gateui' 'gateu' 'gatefCa' 'gatew']

onderzoek(vars, t0, t1, aantal_punten, 20, 'gateh', 'gateu', 'gatef', 'gateu')


    
#    ['u' 'gatef' 'Caup' 'gateui']
#/gnu/store/fk9540xx4zv8bpad7zsvfdln7281ydf3-profile/lib/python3.9/site-packages/pydiffmap/diffusion_map.py:159: RuntimeWarning: invalid value encountered in sqrt
#  dmap = np.dot(evecs, np.diag(np.sqrt(-1. / evals)))
#['gatexr' 'gateh' 'gatefCa' 'Na']
#/gnu/store/fk9540xx4zv8bpad7zsvfdln7281ydf3-profile/lib/python3.9/site-packages/pydiffmap/diffusion_map.py:159: RuntimeWarning: invalid value encountered in sqrt
#  dmap = np.dot(evecs, np.diag(np.sqrt(-1. / evals)))


# ['gatew' 'Carel' 'gatem' 'gatej'], ['gatew' 'gatexs' 'Ca' 'gateu']  --> vertakking
# ['gatej' 'gateoi' 'gateu' 'Caup'] -> loop + ??? vertakking?
# ['gateoi' 'gated' 'gatexs' 'gatej'] --> vertakking

onderzoek(t0, t1, aantal_punten, 0.08, 'gateu', 'Caup', 'gatev', 'gatefCa') # 1D (2D verdwenen?)
onderzoek(t0, t1, aantal_punten, 0.08, 'gatej', 'gateui', 'gatev', 'gateoa') # 2D verdwenen, uitsteeksel??
dmap = onderzoek(t0, t1, aantal_punten, 0.08, 'gateu', 'Caup', 'gateoa', 'gateoi') # 2D verdwenen, vertakking

Another issue is the number of embedding dimensions. The number of significant dimensions of the diffusion map is determined where a remarkable gap occurs in its sorted eigenvalues plot. This is not intrinsically bound to the conventional visualizable dimensions two or three. In contrast, for some other methods such as t-SNE, one can pre-determine the number of visualization dimensions for the embedding optimization to two or three dimensions.
